"""
Stochastic PDE support

This module provides support for stochastic partial differential equations (SPDEs)
with various noise types and integration schemes.
"""

using ..Fields
using ..Grids
using Random
using Statistics

"""
    NoiseProcess

Abstract base type for noise processes in SPDEs.
"""
abstract type NoiseProcess end

"""
    WhiteNoise <: NoiseProcess

Gaussian white noise: η(x,t) with ⟨η(x,t)η(x',t')⟩ = σ²δ(x-x')δ(t-t')
"""
struct WhiteNoise <: NoiseProcess
    variance::Float64
    rng::Random.AbstractRNG
    
    function WhiteNoise(variance=1.0; rng=Random.default_rng())
        new(variance, rng)
    end
end

"""
    ColoredNoise <: NoiseProcess

Spatially correlated noise: η(x,t) with correlation length ξ
"""
struct ColoredNoise <: NoiseProcess
    variance::Float64
    correlation_length::Float64
    rng::Random.AbstractRNG
    
    function ColoredNoise(variance=1.0, correlation_length=1.0; rng=Random.default_rng())
        new(variance, correlation_length, rng)
    end
end

"""
    MultiplicativeNoise <: NoiseProcess

Multiplicative noise: η(x,t)·f(u) where f depends on the field u
"""
struct MultiplicativeNoise <: NoiseProcess
    variance::Float64
    coupling_function::Function
    rng::Random.AbstractRNG
    
    function MultiplicativeNoise(variance=1.0, coupling_func=u->u; rng=Random.default_rng())
        new(variance, coupling_func, rng)
    end
end

"""
    generate_noise(noise::WhiteNoise, field::ScalarField, dt::Float64)

Generate white noise realization for time step dt.
"""
function generate_noise(noise::WhiteNoise, field::ScalarField, dt::Float64)
    field_data = data(field)
    noise_data = sqrt(noise.variance * dt) * randn(noise.rng, size(field_data))
    return ScalarField(noise_data, grid(field))
end

"""
    generate_noise(noise::ColoredNoise, field::ScalarField, dt::Float64)

Generate spatially correlated noise realization.
"""
function generate_noise(noise::ColoredNoise, field::ScalarField, dt::Float64)
    field_data = data(field)
    
    # Generate uncorrelated noise
    white_noise = randn(noise.rng, size(field_data))
    
    # Apply spatial correlation (simple exponential kernel)
    # In practice, this would use FFT convolution for efficiency
    correlated_noise = white_noise  # Simplified - would need proper correlation
    
    noise_data = sqrt(noise.variance * dt) * correlated_noise
    return ScalarField(noise_data, grid(field))
end

"""
    generate_noise(noise::MultiplicativeNoise, field::ScalarField, dt::Float64)

Generate multiplicative noise realization.
"""
function generate_noise(noise::MultiplicativeNoise, field::ScalarField, dt::Float64)
    field_data = data(field)
    coupling = noise.coupling_function(field_data)
    white_noise = randn(noise.rng, size(field_data))
    
    noise_data = sqrt(noise.variance * dt) * coupling .* white_noise
    return ScalarField(noise_data, grid(field))
end

"""
    StochasticPDE <: AbstractPDE

Generic stochastic PDE wrapper that adds noise to deterministic PDEs.

# Mathematical form:
    du/dt = F(u, ∇u, ∇²u, ...) + σ·η(x,t)

where F is the deterministic part and η is a noise process.
"""
struct StochasticPDE <: AbstractPDE
    deterministic_pde::AbstractPDE
    noise_processes::Vector{NoiseProcess}
    noise_coefficients::Vector{Float64}
    
    function StochasticPDE(det_pde::AbstractPDE, noise::NoiseProcess, coefficient=1.0)
        new(det_pde, [noise], [coefficient])
    end
    
    function StochasticPDE(det_pde::AbstractPDE, noises::Vector{NoiseProcess}, 
                          coefficients::Vector{Float64})
        if length(noises) != length(coefficients)
            throw(ArgumentError("Number of noise processes must match coefficients"))
        end
        new(det_pde, noises, coefficients)
    end
end

"""
    evolution_rate(pde::StochasticPDE, state, t::Float64; dt::Float64=0.001)

Compute evolution rate including stochastic terms.
Note: Stochastic term depends on time step dt for proper scaling.
"""
function evolution_rate(pde::StochasticPDE, state, t::Float64; dt::Float64=0.001)
    # Deterministic part
    deterministic_rate = evolution_rate(pde.deterministic_pde, state, t)
    
    # Add noise terms
    stochastic_rate = deterministic_rate
    
    for (i, noise) in enumerate(pde.noise_processes)
        coeff = pde.noise_coefficients[i]
        if isa(state, ScalarField)
            noise_term = coeff * generate_noise(noise, state, dt)
            stochastic_rate = stochastic_rate + noise_term
        elseif isa(state, FieldCollection)
            # Add noise to each field in collection
            noise_fields = []
            for field in state
                noise_field = coeff * generate_noise(noise, field, dt)
                push!(noise_fields, noise_field)
            end
            noise_collection = FieldCollection(noise_fields)
            stochastic_rate = stochastic_rate + noise_collection
        end
    end
    
    return stochastic_rate
end

"""
    StochasticAllenCahnPDE

Allen-Cahn equation with additive noise:
    ∂c/∂t = γ∇²c + c - c³ + √(2T)η(x,t)

where T is the temperature and η is white noise.
"""
function StochasticAllenCahnPDE(; interface_width=1.0, mobility=1.0, temperature=0.1,
                               bc="auto_periodic_neumann", rng=Random.default_rng())
    # Base Allen-Cahn PDE
    det_pde = AllenCahnPDE(interface_width=interface_width, mobility=mobility, bc=bc)
    
    # Thermal noise
    noise = WhiteNoise(2*temperature, rng=rng)
    
    return StochasticPDE(det_pde, noise)
end

"""
    StochasticKPZPDE

KPZ equation - inherently stochastic:
    ∂h/∂t = ν∇²h + (λ/2)|∇h|² + √(2D)η(x,t)

This is already implemented in KPZInterfacePDE but provided here for completeness.
"""
function StochasticKPZPDE(; nu=0.5, lambda=1.0, noise_strength=1.0,
                         bc="auto_periodic_neumann", rng=Random.default_rng())
    return KPZInterfacePDE(nu=nu, lambda=lambda, noise=noise_strength, bc=bc, rng=rng)
end

"""
    StochasticCahnHilliardPDE

Cahn-Hilliard equation with multiplicative noise:
    ∂c/∂t = ∇²[c³ - c - γ∇²c] + ∇²[√(2MT)c(1-c²)η(x,t)]

where the noise is multiplicative and conserves mass on average.
"""
function StochasticCahnHilliardPDE(; interface_width=1.0, mobility=1.0, temperature=0.01,
                                  bc_c="auto_periodic_neumann", bc_mu="auto_periodic_neumann",
                                  rng=Random.default_rng())
    # Base Cahn-Hilliard PDE  
    det_pde = CahnHilliardPDE(interface_width=interface_width, mobility=mobility,
                             bc_c=bc_c, bc_mu=bc_mu)
    
    # Multiplicative noise: coupling ∝ c(1-c²)
    coupling_func = c -> c .* (1.0 .- c.^2)
    noise = MultiplicativeNoise(2*mobility*temperature, coupling_func, rng=rng)
    
    return StochasticPDE(det_pde, noise)
end

"""
    compute_noise_statistics(noise_realizations::Vector{<:AbstractField})

Compute statistics of noise realizations for validation.
"""
function compute_noise_statistics(noise_realizations::Vector{<:AbstractField})
    if isempty(noise_realizations)
        return (mean=0.0, variance=0.0, correlation=0.0)
    end
    
    # Compute mean and variance
    all_data = [data(field) for field in noise_realizations]
    combined_data = vcat([vec(d) for d in all_data]...)
    
    noise_mean = mean(combined_data)
    noise_variance = var(combined_data)
    
    # Compute spatial correlation (simplified)
    correlation_length = 0.0  # Would need proper implementation
    
    return (mean=noise_mean, variance=noise_variance, correlation=correlation_length)
end

"""
    ensemble_average(trajectories::Vector{Vector{<:AbstractField}})

Compute ensemble average over multiple stochastic trajectories.
"""
function ensemble_average(trajectories::Vector{Vector{<:AbstractField}})
    if isempty(trajectories) || isempty(trajectories[1])
        return Float64[]
    end
    
    n_times = length(trajectories[1])
    n_realizations = length(trajectories)
    
    averaged_trajectory = similar(trajectories[1])
    
    for t in 1:n_times
        # Average over all realizations at time t
        field_sum = trajectories[1][t]
        for r in 2:n_realizations
            field_sum = field_sum + trajectories[r][t]
        end
        averaged_trajectory[t] = field_sum / n_realizations
    end
    
    return averaged_trajectory
end

"""
    has_noise(pde::StochasticPDE)

Check if PDE has stochastic terms.
"""
has_noise(pde::StochasticPDE) = !isempty(pde.noise_processes)

"""
    get_noise_strength(pde::StochasticPDE)

Get total noise strength.
"""
function get_noise_strength(pde::StochasticPDE)
    total_strength = 0.0
    for (i, noise) in enumerate(pde.noise_processes)
        coeff = pde.noise_coefficients[i]
        total_strength += coeff^2 * noise.variance
    end
    return sqrt(total_strength)
end

# Export stochastic PDE framework
export NoiseProcess, WhiteNoise, ColoredNoise, MultiplicativeNoise
export generate_noise, StochasticPDE
export StochasticAllenCahnPDE, StochasticKPZPDE, StochasticCahnHilliardPDE
export compute_noise_statistics, ensemble_average, has_noise, get_noise_strength